package com.niodata.md.rdb.crawler;

import com.niodata.md.rdb.crawler.plugin.DorisRowCountRetrieval;
import com.niodata.md.rdb.crawler.plugin.DorisTableStorageSizeRetrieval;
import com.niodata.md.rdb.model.MetaType;
import java.util.HashMap;
import java.util.Map;

/**
 * metadata retrieval registry.
 */
public class MetadataRetrievalRegistry {

  //dbcategory.metaType -> retrievalClass
  private static final Map<String, Class<? extends AbstractRetrieval>> map = new HashMap<>();

  static {
    map.put("doris." + MetaType.TABLE_SIZE, DorisTableStorageSizeRetrieval.class);
    map.put("doris." + MetaType.ROW_COUNT, DorisRowCountRetrieval.class);
  }

  /**
   * register retrieval class.
   *
   * @param dbCategory db category
   * @param metaType   meta type
   * @param clazz      class to register
   */
  public static void registerRetrieval(String dbCategory, MetaType metaType, Class<? extends AbstractRetrieval> clazz) {
    String key = dbCategory + "." + metaType;
    map.put(key, clazz);
  }


  /**
   * get retrieval class.
   *
   * @param dbCategory db category
   * @param metaType   meta type
   * @return retrieval class
   */
  public static Class<? extends AbstractRetrieval> getRetrievalClass(String dbCategory, MetaType metaType) {
    String key = dbCategory + "." + metaType;
    return map.get(key);
  }

}
